{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from ipywidgets import interactive\n",
    "import ipywidgets as widgets\n",
    "from ipywidgets import fixed\n",
    "from sklearn.linear_model import LinearRegression\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_float_widget(name):\n",
    "    return widgets.FloatLogSlider(\n",
    "        value=-10,\n",
    "        base=np.sqrt(2),\n",
    "        min=-30,\n",
    "        max=0,\n",
    "        step=1,\n",
    "        description=name+':',\n",
    "        continuous_update= False\n",
    "    )\n",
    "\n",
    "def generate_int_widget(name, val, min_, max_, step=1):\n",
    "    return widgets.IntSlider(\n",
    "        value=val,\n",
    "        min=min_,\n",
    "        max=max_,\n",
    "        step=step,\n",
    "        description=name + ':',\n",
    "        disabled=False,\n",
    "        continuous_update=False,\n",
    "        orientation='horizontal',\n",
    "        readout=True,\n",
    "        readout_format='d')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_true_weights(k, p):\n",
    "    w_true = np.zeros(k+p)\n",
    "    w_true[:k] = np.random.randn(k)\n",
    "    #we put the true weights on the sphere, so ||w^*||=1\n",
    "    w_true /= np.linalg.norm(w_true)\n",
    "    return w_true\n",
    "\n",
    "\n",
    "def generate_X_orthogonal_blocks(lambda_L, lambda_S, k, n, p):\n",
    "    X = np.empty(shape=(n, k+p))\n",
    "    X[:, :k] = np.linalg.svd(np.random.randn(n, k), full_matrices=False)[0] * np.sqrt(lambda_L * n)\n",
    "    X[:, k:] = np.linalg.svd(np.random.randn(n,p), full_matrices=False)[2] * np.sqrt(lambda_S * p)\n",
    "    return X\n",
    "def generate_X_gaussian(lambda_L, lambda_S, k, n, p):\n",
    "    X = np.empty(shape=(n, k+p))\n",
    "    X[:, :k] = np.random.randn(n,k) * np.sqrt(lambda_L)\n",
    "    X[:, k:] = np.random.randn(n,p) * np.sqrt(lambda_S)\n",
    "    return X\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 1\n",
    "In this part of the notebook we will check that computations from part e of the problem. Below there are four cells with the code. Each of those cells contains a function that generates the data and computes one of four terms of the error from part b. There is also the second function that is not implemented yet. **Add the code to the second function in each cell so that it would compute the formula for the corresponding error term that you found in part e.** When you run one of those cells, you will be able to choose the parameters and print outputs of both functions for those parameters. If you do everything correctly, both outputs will be exactly the same!\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The next cell computes the error term that corresponds to **bias in the spiked part.** We take  $\\|{\\bf w}^*\\|$ to be 1 as the error term doesn't have any interesting dependence on this quantity (they are simply proportional). We also set $\\lambda_L = 1$ and $n = 100$ and allow changing $k,p,\\lambda_S$.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_value = 100\n",
    "\n",
    "def bias_spiked_empirical(lambda_L, lambda_S, k, n, p):\n",
    "    lambdas = np.array([lambda_L] * k + [lambda_S] * p)\n",
    "    X = generate_X_orthogonal_blocks(lambda_L, lambda_S, k, n, p)\n",
    "    w_true = generate_true_weights(k, p)\n",
    "    w_estimated = X.T @ np.linalg.inv(X @ X.T) @ X @ w_true\n",
    "    return np.sum((lambdas * (w_true - w_estimated)**2)[:k])\n",
    "\n",
    "def bias_spiked_theoretical(lambda_L, lambda_S, k, n, p):\n",
    "    #TODO implement this function\n",
    "    ### start bias_spiked ###\n",
    "\n",
    "    ### end bias_spiked ###\n",
    "\n",
    "def compare_bias_spiked(lambda_L, lambda_S, k, n, p):\n",
    "    print(\"simulation result: \", bias_spiked_empirical(lambda_L, lambda_S, k, n, p))\n",
    "    print(\"theoretical result: \", bias_spiked_theoretical(lambda_L, lambda_S, k, n, p))\n",
    "\n",
    "interactive_comparison = interactive(compare_bias_spiked,\n",
    "                                     lambda_L=fixed(1.),\n",
    "                                     lambda_S=generate_float_widget('$\\lambda_S$'),\n",
    "                                     k=generate_int_widget('k', n_value//2, 1, n_value),\n",
    "                                     n=fixed(n_value),\n",
    "                                     p=generate_int_widget('p', 2*n_value, n_value, 10 * n_value))\n",
    "\n",
    "interactive_comparison\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The next cell computes the error term that corresponds to **bias in the non-spiked part.** We take  $\\|{\\bf w}^*\\|$ to be 1 as the error term doesn't have any interesting dependence on this quantity (they are simply proportional). We also set $\\lambda_L = 1$ and $n = 100$ and allow changing $k,p,\\lambda_S$.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_value = 100\n",
    "\n",
    "def bias_non_spiked_empirical(lambda_L, lambda_S, k, n, p):\n",
    "    lambdas = np.array([lambda_L] * k + [lambda_S] * p)\n",
    "    X = generate_X_orthogonal_blocks(lambda_L, lambda_S, k, n, p)\n",
    "    w_true = generate_true_weights(k, p)\n",
    "    w_estimated = X.T @ np.linalg.inv(X @ X.T) @ X @ w_true\n",
    "    return np.sum((lambdas * (w_true - w_estimated)**2)[k:])\n",
    "\n",
    "def bias_non_spiked_theoretical(lambda_L, lambda_S, k, n, p):\n",
    "    #TODO implement this function\n",
    "    ### start bias_non-spiked ###\n",
    "\n",
    "    ### end bias_non-spiked ###\n",
    "\n",
    "def compare_bias_non_spiked(lambda_L, lambda_S, k, n, p):\n",
    "    print(bias_non_spiked_empirical(lambda_L, lambda_S, k, n, p))\n",
    "    print(bias_non_spiked_theoretical(lambda_L, lambda_S, k, n, p))\n",
    "\n",
    "interactive_comparison = interactive(compare_bias_non_spiked,\n",
    "                                     lambda_L=fixed(1.),\n",
    "                                     lambda_S=generate_float_widget('$\\lambda_S$'),\n",
    "                                     k=generate_int_widget('k', n_value//2, 1, n_value),\n",
    "                                     n=fixed(n_value),\n",
    "                                     p=generate_int_widget('p', 2*n_value, n_value, 10 * n_value))\n",
    "\n",
    "interactive_comparison\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The next cell computes the error term that corresponds to **variance in the spiked part.** We take  $\\|\\sigma\\|$ to be 1 as the error term doesn't have any interesting dependence on this quantity (they are simply proportional). We also set $\\lambda_L = 1$ and $n = 100$ and allow changing $k,p,\\lambda_S$.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_value = 100\n",
    "\n",
    "def variance_spiked_empirical(lambda_L, lambda_S, k, n, p):\n",
    "    lambdas = np.array([lambda_L] * k + [lambda_S] * p)\n",
    "    X = generate_X_orthogonal_blocks(lambda_L, lambda_S, k, n, p)\n",
    "    return np.linalg.norm((np.diag(np.sqrt(lambdas)) @ X.T @ np.linalg.inv(X @ X.T))[:k, :])**2\n",
    "\n",
    "def variance_spiked_theoretical(lambda_L, lambda_S, k, n, p):\n",
    "    #TODO implement this function\n",
    "    ### start variance_spiked ###\n",
    "\n",
    "    ### end variance_spiked ###\n",
    "\n",
    "def compare_variance_spiked(lambda_L, lambda_S, k, n, p):\n",
    "    print(variance_spiked_empirical(lambda_L, lambda_S, k, n, p))\n",
    "    print(variance_spiked_theoretical(lambda_L, lambda_S, k, n, p))\n",
    "\n",
    "\n",
    "interactive_comparison = interactive(compare_variance_spiked,\n",
    "                                     lambda_L=fixed(1.),\n",
    "                                     lambda_S=generate_float_widget('$\\lambda_S$'),\n",
    "                                     k=generate_int_widget('k', n_value//2, 1, n_value),\n",
    "                                     n=fixed(n_value),\n",
    "                                     p=generate_int_widget('p', 2*n_value, n_value, 10 * n_value))\n",
    "\n",
    "interactive_comparison\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The next cell computes the error term that corresponds to **variance in the non-spiked part.** We take  $\\|\\sigma\\|$ to be 1 as the error term doesn't have any interesting dependence on this quantity (they are simply proportional). We also set $\\lambda_L = 1$ and $n = 100$ and allow changing $k,p,\\lambda_S$.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_value = 100\n",
    "\n",
    "def variance_non_spiked_empirical(lambda_L, lambda_S, k, n, p):\n",
    "    lambdas = np.array([lambda_L] * k + [lambda_S] * p)\n",
    "    X = generate_X_orthogonal_blocks(lambda_L, lambda_S, k, n, p)\n",
    "    return np.linalg.norm((np.diag(np.sqrt(lambdas)) @ X.T @ np.linalg.inv(X @ X.T))[k:, :])**2\n",
    "\n",
    "def variance_non_spiked_theoretical(lambda_L, lambda_S, k, n, p):\n",
    "    #TODO implement this function\n",
    "    ### start variance_non-spiked ###\n",
    "\n",
    "    ### end variance_non-spiked ###\n",
    "\n",
    "def compare_variance_non_spiked(lambda_L, lambda_S, k, n, p):\n",
    "    print(variance_non_spiked_empirical(lambda_L, lambda_S, k, n, p))\n",
    "    print(variance_non_spiked_theoretical(lambda_L, lambda_S, k, n, p))\n",
    "\n",
    "interactive_comparison = interactive(compare_variance_non_spiked,\n",
    "                                     lambda_L=fixed(1.),\n",
    "                                     lambda_S=generate_float_widget('$\\lambda_S$'),\n",
    "                                     k=generate_int_widget('k', n_value//2, 1, n_value),\n",
    "                                     n=fixed(n_value),\n",
    "                                     p=generate_int_widget('p', 2*n_value, n_value, 10 * n_value))\n",
    "\n",
    "interactive_comparison\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 2\n",
    "Recall that in part c we introduced assumptions on ${\\bf X}$ that we thought would be approximately satisfied by gaussian design. In this part we are goin to see if it is indeed true.\n",
    "\n",
    "**For $n \\in \\{10, 30, 100\\}$ find such minimum $d$ that if you generate an $n\\times d$ matrix ${\\bf Z}$ with i.i.d. standard normal entries, then it often holds**\n",
    "$$\n",
    "\\|{\\bf Z}{\\bf Z}^\\top/d - {\\bf I}_n\\|< 1/3\n",
    "$$\n",
    "\n",
    "**How big do you think  $p/k$ should be to make gaussian design behave similarly to the design that we studied in our problem?**\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### start check approximate conditions ###\n",
    "n = 10\n",
    "d = 400\n",
    "Z = np.random.randn(n, d)\n",
    "sing_vals = np.linalg.svd(Z)[1]\n",
    "print(max(np.abs(sing_vals[0]**2 - d), np.abs(sing_vals[-1]**2 - d))/d)\n",
    "### end check approximate conditions ###\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 3\n",
    "\n",
    "Finally, we check how well our bounds work  for the gaussian design. Below there are four cells with code. In each of them we compare the error terms that we see for gaussian design with the bounds that we obtained for our simplified design. As in part a, we put $n=100$, $\\lambda_L = 1$, $\\|{\\bf w^*}\\|=1$ and $\\sigma = 1$. Then for different values of $k$ and $d$ we compare our bounds for the 4 error terms to the corresponding error terms for gaussian design. **Run the experiements in this part and answer the following:**\n",
    "1. **Do bounds from part e still work reasonably well if we have gaussian design?**\n",
    "2. **For which values of the parameters the discrepancy is the highest?**\n",
    "3. **Was your guess about $p/k$ from the previous part correct?**\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Bias in the spiked part:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_value = 100\n",
    "\n",
    "def bias_spiked_gaussian(lambda_L, lambda_S, k, n, p):\n",
    "    lambdas = np.array([lambda_L] * k + [lambda_S] * p)\n",
    "    X = generate_X_gaussian(lambda_L, lambda_S, k, n, p)\n",
    "    w_true = generate_true_weights(k, p)\n",
    "    w_estimated = X.T @ np.linalg.inv(X @ X.T) @ X @ w_true\n",
    "    return np.sum((lambdas * (w_true - w_estimated)**2)[:k])\n",
    "\n",
    "\n",
    "def compare_with_gaussian_bias_spiked(lambda_L, lambda_S_range, k, n, p):\n",
    "    gaussian_errors=[]\n",
    "    our_bounds=[]\n",
    "    for lambda_S in lambda_S_range:\n",
    "        gaussian_errors.append(bias_spiked_gaussian(lambda_L, lambda_S, k, n, p))\n",
    "        our_bounds.append(bias_spiked_theoretical(lambda_L, lambda_S, k, n, p))\n",
    "    plt.plot(lambda_S_range, gaussian_errors, label='gaussian X')\n",
    "    plt.plot(lambda_S_range, our_bounds, label='our bound')\n",
    "    plt.ylim([0,1])\n",
    "    plt.xlabel(\"$\\lambda_S$\")\n",
    "    plt.ylabel(\"error\")\n",
    "    plt.legend()\n",
    "\n",
    "\n",
    "interactive_comparison = interactive(compare_with_gaussian_bias_spiked,\n",
    "                                     lambda_L=fixed(1.),\n",
    "                                     lambda_S_range=fixed(np.linspace(0.01, 1, 100)),\n",
    "                                     k=generate_int_widget('k', n_value//2, 1, n_value),\n",
    "                                     n=fixed(n_value),\n",
    "                                     p=generate_int_widget('p', 2*n_value, n_value, 10 * n_value))\n",
    "\n",
    "interactive_comparison\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Bias in the non-spiked part:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_value = 100\n",
    "\n",
    "def bias_non_spiked_gaussian(lambda_L, lambda_S, k, n, p):\n",
    "    lambdas = np.array([lambda_L] * k + [lambda_S] * p)\n",
    "    X = generate_X_gaussian(lambda_L, lambda_S, k, n, p)\n",
    "    w_true = generate_true_weights(k, p)\n",
    "    w_estimated = X.T @ np.linalg.inv(X @ X.T) @ X @ w_true\n",
    "    return np.sum((lambdas * (w_true - w_estimated)**2)[k:])\n",
    "\n",
    "\n",
    "def compare_with_gaussian_bias_non_spiked(lambda_L, lambda_S_range, k, n, p):\n",
    "    gaussian_errors=[]\n",
    "    our_bounds=[]\n",
    "    for lambda_S in lambda_S_range:\n",
    "        gaussian_errors.append(bias_non_spiked_gaussian(lambda_L, lambda_S, k, n, p))\n",
    "        our_bounds.append(bias_non_spiked_theoretical(lambda_L, lambda_S, k, n, p))\n",
    "    plt.plot(lambda_S_range, gaussian_errors, label='gaussian X')\n",
    "    plt.plot(lambda_S_range, our_bounds, label='our bound')\n",
    "    plt.ylim([0,1])\n",
    "    plt.xlabel(\"$\\lambda_S$\")\n",
    "    plt.ylabel(\"error\")\n",
    "    plt.legend()\n",
    "\n",
    "interactive_comparison = interactive(compare_with_gaussian_bias_non_spiked,\n",
    "                                     lambda_L=fixed(1.),\n",
    "                                     lambda_S_range=fixed(np.linspace(0.1, 1, 100)),\n",
    "                                     k=generate_int_widget('k', n_value//2, 1, n_value),\n",
    "                                     n=fixed(n_value),\n",
    "                                     p=generate_int_widget('p', 2*n_value, n_value, 10 * n_value))\n",
    "\n",
    "interactive_comparison\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Variance in the spiked part:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_value = 100\n",
    "\n",
    "def variance_spiked_gaussian(lambda_L, lambda_S, k, n, p):\n",
    "    lambdas = np.array([lambda_L] * k + [lambda_S] * p)\n",
    "    X = generate_X_gaussian(lambda_L, lambda_S, k, n, p)\n",
    "    return np.linalg.norm((np.diag(np.sqrt(lambdas)) @ X.T @ np.linalg.inv(X @ X.T))[:k, :])**2\n",
    "\n",
    "def compare_with_gaussian_variance_spiked(lambda_L, lambda_S_range, k, n, p):\n",
    "    gaussian_errors=[]\n",
    "    our_bounds=[]\n",
    "    for lambda_S in lambda_S_range:\n",
    "        gaussian_errors.append(variance_spiked_gaussian(lambda_L, lambda_S, k, n, p))\n",
    "        our_bounds.append(variance_spiked_theoretical(lambda_L, lambda_S, k, n, p))\n",
    "    plt.plot(lambda_S_range, gaussian_errors, label='gaussian X')\n",
    "    plt.plot(lambda_S_range, our_bounds, label='our bound')\n",
    "    plt.ylim([0,1])\n",
    "    plt.xlabel(\"$\\lambda_S$\")\n",
    "    plt.ylabel(\"error\")\n",
    "    plt.legend()\n",
    "\n",
    "interactive_comparison = interactive(compare_with_gaussian_variance_spiked,\n",
    "                                     lambda_L=fixed(1.),\n",
    "                                     lambda_S_range=fixed(np.linspace(0.1, 1, 100)),\n",
    "                                     k=generate_int_widget('k', n_value//2, 1, n_value),\n",
    "                                     n=fixed(n_value),\n",
    "                                     p=generate_int_widget('p', 2*n_value, n_value, 10 * n_value))\n",
    "\n",
    "interactive_comparison\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Variance in the non-spiked part:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_value = 100\n",
    "\n",
    "def variance_non_spiked_gaussian(lambda_L, lambda_S, k, n, p):\n",
    "    lambdas = np.array([lambda_L] * k + [lambda_S] * p)\n",
    "    X = generate_X_gaussian(lambda_L, lambda_S, k, n, p)\n",
    "    return np.linalg.norm((np.diag(np.sqrt(lambdas)) @ X.T @ np.linalg.inv(X @ X.T))[k:, :])**2\n",
    "\n",
    "\n",
    "def compare_with_gaussian_variance_non_spiked(lambda_L, lambda_S_range, k, n, p):\n",
    "    gaussian_errors=[]\n",
    "    our_bounds=[]\n",
    "    for lambda_S in lambda_S_range:\n",
    "        gaussian_errors.append(variance_non_spiked_gaussian(lambda_L, lambda_S, k, n, p))\n",
    "        our_bounds.append(variance_non_spiked_theoretical(lambda_L, lambda_S, k, n, p))\n",
    "    plt.plot(lambda_S_range, gaussian_errors, label='gaussian X')\n",
    "    plt.plot(lambda_S_range, our_bounds, label='our bound')\n",
    "    plt.ylim([0,1])\n",
    "    plt.xlabel(\"$\\lambda_S$\")\n",
    "    plt.ylabel(\"error\")\n",
    "    plt.legend()\n",
    "\n",
    "interactive_comparison = interactive(compare_with_gaussian_variance_non_spiked,\n",
    "                                     lambda_L=fixed(1.),\n",
    "                                     lambda_S_range=fixed(np.linspace(0.1, 1, 100)),\n",
    "                                     k=generate_int_widget('k', n_value//2, 1, n_value),\n",
    "                                     n=fixed(n_value),\n",
    "                                     p=generate_int_widget('p', 2*n_value, n_value, 10 * n_value))\n",
    "\n",
    "interactive_comparison\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
